from telnetlib import Telnet
import re

HOST = "localhost"
PORT = 4444
TIMEOUT = 10

class telnet_openocd:
    def __init__(self):
        """连接到 OpenOCD 的 Telnet 接口"""
        try:
            self.tn = Telnet(HOST, PORT, TIMEOUT)
            print(f"成功连接到 {HOST}:{PORT}")
            print(self.tn.read_until(b'>', TIMEOUT).decode('ascii'))
        except Exception as e:
            print(f"无法连接到 OpenOCD: {e}")
            exit()

    def send_command(self, command):
        """发送命令到 OpenOCD 并返回结果"""
        self.tn.write(command.encode('ascii') + b'\n')
        result = self.tn.read_until(b'>', TIMEOUT).decode('ascii')
        return result

    def read_pc(self):
        """读取 PC 寄存器的值"""
        result = self.send_command("reg pc")
        lines = result.split('\n')
        for line in lines:
            if 'pc' in line:
                parts = line.split(':')
                if len(parts) > 1:
                    pc_value = parts[1].strip()
                    # print(f"PC 寄存器的值: {pc_value}")
                    return pc_value
        # print("未找到 PC 寄存器的值")
        return None

    def write_pc(self, value):
        """设置 PC 寄存器的值"""
        command = f"reg pc {value}"
        result = self.send_command(command)
        # print(f"设置 PC 寄存器的结果: {result}")

    def read_ra(self):
        """读取 ra 寄存器的值"""
        result = self.send_command("reg ra")
        lines = result.split('\n')
        for line in lines:
            if 'ra' in line:
                parts = line.split(':')
                if len(parts) > 1:
                    ra_value = parts[1].strip()
                    # print(f"ra 寄存器的值: {ra_value}")
                    return ra_value
        # print("未找到 ra 寄存器的值")
        return None

    def write_ra(self, value):
        """设置 ra 寄存器的值"""
        command = f"reg ra {value}"
        result = self.send_command(command)
        # print(f"设置 ra 寄存器的结果: {result}")

    def read_sp(self):
        """读取 sp 寄存器的值"""
        result = self.send_command("reg sp")
        lines = result.split('\n')
        for line in lines:
            if 'sp' in line:
                parts = line.split(':')
                if len(parts) > 1:
                    sp_value = parts[1].strip()
                    # print(f"sp 寄存器的值: {sp_value}")
                    return sp_value
        # print("未找到 sp 寄存器的值")
        return None

    def write_sp(self, value):
        """设置 sp 寄存器的值"""
        command = f"reg sp {value}"
        result = self.send_command(command)
        # print(f"设置 sp 寄存器的结果: {result}")

    def read_memory_byte(self, address):
        cmd = f"mdb 0x{address:2X}"
        
        response = self.send_command(cmd)
        
        # 解析响应
        try:
            match = re.search(r':\s*([^\r\n]+)', response)
            if match:
                value = match.group(1).strip()
            return int(value, 16)
        except Exception as e:
            print(f"Error parsing response: {response}")
            raise e

    def write_memory_byte(self, address, value):
        cmd = f"mwb 0x{address:02X} {value:02X}"
        self.send_command(cmd)

    def reset(self):
        result = self.send_command("reset")

    def reset_run(self):
        result = self.send_command("reset run")

    def reset_init(self):
        result = self.send_command("reset init")

    def reset_halt(self):
        result = self.send_command("reset halt")

    def resume(self):
        result = self.send_command("resume")

    def halt(self):
        result = self.send_command("halt")

    def set_breakpoint(self, address):
        cmd = f"bp 0x{address:02X} 4 hw"
        self.send_command(cmd)
